#include <iostream>
#include <fstream>
#include <cstring>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <chrono>

#define PORT 8081
#define BUFFER_SIZE 65507  // UDP最大数据报大小
#define MAX_ATTEMPTS 5

int main(int argc, char const *argv[]) {
    int sockfd;
    struct sockaddr_in servaddr;
    char buffer[BUFFER_SIZE] = {0};
    const char* filename = "large_file.bin";
    
    if (argc > 1) {
        filename = argv[1];
    }
    
    // 创建套接字
    if ((sockfd = socket(AF_INET, SOCK_DGRAM, 0)) < 0) {
        perror("socket creation failed");
        exit(EXIT_FAILURE);
    }
    
    memset(&servaddr, 0, sizeof(servaddr));
    
    // 设置服务器地址
    servaddr.sin_family = AF_INET;
    servaddr.sin_port = htons(PORT);
    servaddr.sin_addr.s_addr = INADDR_ANY;
    
    std::ifstream file(filename, std::ios::binary);
    if (!file) {
        perror("file open failed");
        exit(EXIT_FAILURE);
    }
    
    auto start_time = std::chrono::high_resolution_clock::now();
    size_t total_bytes = 0;
    socklen_t len = sizeof(servaddr);
    
    // 发送文件
    while (!file.eof()) {
        file.read(buffer, BUFFER_SIZE);
        size_t bytes_to_send = file.gcount();
        if (bytes_to_send > 0) {
            // 可靠UDP：发送数据并等待确认
            int attempts = 0;
            bool acknowledged = false;
            
            while (attempts < MAX_ATTEMPTS && !acknowledged) {
                sendto(sockfd, (const char *)buffer, bytes_to_send, 
                       MSG_CONFIRM, (const struct sockaddr *)&servaddr, len);
                
                // 等待确认
                fd_set readfds;
                struct timeval timeout;
                timeout.tv_sec = 1;  // 1秒超时
                timeout.tv_usec = 0;
                
                FD_ZERO(&readfds);
                FD_SET(sockfd, &readfds);
                
                int activity = select(sockfd + 1, &readfds, NULL, NULL, &timeout);
                if (activity > 0) {
                    char ack_buffer[10];
                    recvfrom(sockfd, (char *)ack_buffer, 10, 
                            MSG_WAITALL, (struct sockaddr *)&servaddr, &len);
                    if (std::string(ack_buffer) == "ACK") {
                        acknowledged = true;
                    }
                }
                attempts++;
            }
            
            total_bytes += bytes_to_send;
        }
    }
    
    // 发送结束标记
    sendto(sockfd, "", 0, MSG_CONFIRM, 
           (const struct sockaddr *)&servaddr, len);
    
    auto end_time = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::seconds>(end_time - start_time).count();
    
    file.close();
    close(sockfd);
    
    std::cout << "UDP File sent successfully: " << total_bytes << " bytes in " 
              << duration << " seconds" << std::endl;
    if (duration > 0) {
        std::cout << "Effective transfer rate: " 
                  << (total_bytes / (1024.0 * 1024.0)) / duration << " MB/s" << std::endl;
    }
    
    return 0;
}    